"""
自动生成推荐女声音色
"""

import torch
import ChatTTS
import torchaudio
import numpy as np
from pathlib import Path


def main():
    print("=" * 70)
    print("ChatTTS 女声音色自动生成工具")
    print("=" * 70)
    
    # 初始化ChatTTS
    print("\n[1/3] 正在初始化ChatTTS引擎...")
    chat = ChatTTS.Chat()
    chat.load(compile=False, source='huggingface')
    print("✓ ChatTTS引擎初始化完成")
    
    # 创建音色保存目录
    speaker_dir = Path("speakers")
    speaker_dir.mkdir(exist_ok=True)
    
    # 推荐的女声种子
    female_seeds = {
        2222: "温柔女声",
        6666: "活泼女声",
        9527: "甜美女声",
        1234: "清新女声",
        8888: "成熟女声"
    }
    
    test_text = "大家好,我是AI虚拟主播小爱!"
    
    print(f"\n[2/3] 正在生成 {len(female_seeds)} 个推荐女声...")
    print(f"测试文本: {test_text}")
    print("-" * 70)
    
    for seed, description in female_seeds.items():
        print(f"\n音色种子: {seed} ({description})")
        
        # 设置随机种子
        torch.manual_seed(seed)
        spk_emb = chat.sample_random_speaker()
        
        # 合成测试音频
        params_infer_code = ChatTTS.Chat.InferCodeParams(
            spk_emb=spk_emb,
            temperature=0.3,
            top_P=0.7,
            top_K=20,
        )
        
        wavs = chat.infer([test_text], params_infer_code=params_infer_code)
        
        if wavs is not None and len(wavs) > 0:
            # 保存试听文件
            preview_file = f"female_preview_{seed}.wav"
            try:
                audio_tensor = torch.from_numpy(wavs[0]).unsqueeze(0)
                torchaudio.save(preview_file, audio_tensor, 24000)
            except:
                # 备用方法
                import wave
                audio_int16 = (wavs[0] * 32767).astype(np.int16)
                with wave.open(preview_file, 'wb') as wf:
                    wf.setnchannels(1)
                    wf.setsampwidth(2)
                    wf.setframerate(24000)
                    wf.writeframes(audio_int16.tobytes())
            
            # 保存音色文件
            speaker_file = speaker_dir / f"female_{seed}.pt"
            torch.save(spk_emb, speaker_file)
            
            print(f"  ✓ 试听文件: {preview_file}")
            print(f"  ✓ 音色文件: {speaker_file}")
        else:
            print(f"  ✗ 生成失败")
    
    print("\n" + "=" * 70)
    print("[3/3] 女声音色生成完成!")
    print("=" * 70)
    print("\n请试听以下文件选择你喜欢的音色:")
    for seed, description in female_seeds.items():
        print(f"  - female_preview_{seed}.wav ({description})")
    
    print("\n音色文件已保存到 speakers/ 目录:")
    for seed in female_seeds.keys():
        print(f"  - speakers/female_{seed}.pt")
    
    print("\n" + "=" * 70)
    print("如何使用:")
    print("=" * 70)
    print("\n1. 试听所有音频文件,找到喜欢的音色")
    print("\n2. 编辑 config/config.yaml,设置音色:")
    print("\n   tts:")
    print("     engine: \"chattts\"")
    print("     chattts:")
    print("       speaker_file: \"speakers/female_2222.pt\"  # 选择你喜欢的")
    print("\n3. 运行测试:")
    print("   python test_local_enhanced.py")
    print("\n" + "=" * 70)


if __name__ == "__main__":
    main()
